Sparse linear algebraic methods and MPI with spatio-temporal models

Environmental systems are characterised by phenomena which evolve both in space and in time (weather and climate, glacial flow, spread of invasive species etc.).

The aim of spatio-temporal modelling is to

• Assimilate all knowledge (observables etc.) pertaining to the process/es of interest.
• Infer parameters governing the spatio-temporal process in order to
  • Further our understanding of the process and
  • Predict in both space and time.

• Geo-statistical models:

\[ X_t(\svec) \sim GP(\mu_t(\svec),K(t,t',\svec,\svec')) \]

General-purpose model, used when the dynamic structure is largely unknown.

• Dynamic (mechanistically-motivated) models

\[ X_t(\svec) = f(X_{t-1}(\svec), X_{t-2}(\svec), \dots, \wvec_t(\svec); \thetab_m) \]

• Examples:
  • STAR models
  • Integro-difference equation models
  • PDE models
  • CMLs

• Used when we have some mechanistic insight of the underlying process.
• In the environmental sciences these are usually in the form of (stochastic) PDEs.
• They implicitly define very complicated spatio-temporal covariances.

\[ \frac{\partial X(s;t)}{\partial t} - \beta \frac{\partial^2X(s;t)}{\partial s^2} + a X(s;t) = e(s;t) \]

• Continuous-time ST processes are computationally tedious to deal with.

• Focus on models of the form

\[ X_t(\svec) = f(X_{t-1}(\svec)) + e_t(\svec) \]

where \( ~e_t(\svec)~ \) is the solution to

\[ (\kappa^2 - \Delta)^{\alpha/2}(\tau e(\svec)) = \mathcal{W(\svec)} \]

\[ (\kappa^2 - \Delta)^{\alpha/2}(\tau e(\svec)) = \mathcal{W(\svec)} \]

has as solution (Whittle, 1954)

\[ e(\svec) \sim \mathcal{GP}(0,K(\svec,\svec'; \alpha,\tau,\kappa)) \]

where \( k(\cdot,\cdot) \) is a Matérn covariance function. The wave spectrum of \( e(\svec) \) is given by

\[ R(\kvec) = (2\pi)^{-\alpha}(\kappa^2 + \| \kvec \|^2)^{-\alpha} \]

1. Non-stationarity can be handled by making the SPDE spatially varying. This results in a covariance function which is not a Matérn field, but do we care?
   

2. We can model SPDEs on manifolds, not restricted to \( \mathbb{R}^2 \).
   

3. SPDEs can be approximated to arbitrary precision using an appropriate basis (Lindgren, 2011).

• Useful when we have multiple interacting spatio-temporal processes.

• Consider the bi-variate sub-system of random forcings: \( e_{1,t}(\svec), e_{2,t}(\svec) \). Define the multi-variate Gaussian process (MVGP) as

\[ (e_{1,t}(\svec), e_{2,t}(\svec)) = \bold{e}(\svec) \sim \mathcal{MVGP}(\bold{0},\bold{K}(\svec,\svec',t,t';\thetab_K)) \]

where \( \bold{K}(\svec,\svec',t,t';\thetab_K) \) is a \( 2 \times 2 \) cross-covariance matrix function (i.e.~ each element of \( \bold{K} \) is a covariance function (Finley et al., 2007):

\[ \bold{K}(\svec,\svec',t,t') = [Cov(e_{i,t}(\svec), e_{i,t'}(\svec'))]_{i,j \in \{1,2\}} \]

• Model with a system of SPDEs. Example of a three-variate system of random forcings:

\[ \begin{bmatrix} \mathcal{A}_1 & & \\ & & \mathcal{A}_2 & \mathcal{A}_{2,3} \\ & & \mathcal{A}_{3,2} & \mathcal{A}_3 \end{bmatrix} \begin{bmatrix} e_1(\svec) \\ e_2(\svec) \\ e_3(\svec) \end{bmatrix} = \begin{bmatrix} \mathcal{W}_1(\svec) \\ \mathcal{W}_2(\svec) \\ \mathcal{W}_3(\svec) \end{bmatrix} \]

• SPDEs are a useful tool in modelling complex phenomena.

• Real power comes in using finite elements to decompose and represent them as GMRFs.

• From now on we restrict ourselves to the linear, Gaussian case

\[ \begin{bmatrix} \xvec_{1,t} \\ \xvec_{2,t} \\ \vdots \end{bmatrix} = \begin{bmatrix} \Amat_{1} & \Amat_{1,2} & \dots \\ \Amat_{2,1}& \Amat_{2,2} & \dots \\ \vdots& \vdots & \ddots \\ \end{bmatrix} \begin{bmatrix} \xvec_{1,t-1} \\ \xvec_{2,t-1} \\ \vdots \end{bmatrix} + \bold{J}_t\bold{\beta} + \begin{bmatrix} \evec_{1,t} \\ \evec_{2,t} \\ \vdots \end{bmatrix} \]

• Kalman filter/ RTS methods do not lend themselves to sparse methods (more of that in another talk).

• Collapse all variates, and their spatial and temporal components into one big GMRF.

• Sparsity is retained and one can make use of fill-in reduction permutations when computing the Cholesky factor (Knorr-Rue and Held, 2002)

\[ p(\xvec,\betab_I | \Thetab) = \left(\prod_{i=1}^T{p(\xvec_{t} | \xvec_{t-1},\betab,\Thetab)}\right)p(\xvec_0 | \betab, \Thetab)p(\betab_I) \]

• Initial distribution (Note, not the stationary one!):

\[ \xvec_0|\betab_I,\Thetab \sim \mathcal{N}(\bold{J}_0\betab,[\Qmat_w - \Amat^T\Qmat_w\Amat]^{-1}) \]

\[ \small \begin{array}{cc} \Qmat = \begin{bmatrix} \Qmat_1 & \Qmat_2 \\ \Qmat_2^T & \Qmat_3 \end{bmatrix} & \Qmat_2 = \begin{bmatrix} \Amat^T\Qmat_w\Jmat_{1} - \Qmat_0\Jmat_{0} \\ \Amat^T\Qmat_w\Jmat_{2} - \Qmat_w\Jmat_1 \\ \vdots\end{bmatrix}, \end{array} \]

\[ \small \Qmat_1 = \begin{bmatrix} \Amat^T\Qmat_w\Amat + \Qmat_0 & -\Amat^T\Qmat_w & ~~~ \\ -\Qmat_w\Amat & \Amat^T\Qmat_w\Amat + \Qmat_w& -\Amat^T\Qmat_w & ~~~ \\ & -\Qmat_w\Amat & \ddots & \ddots \\ & \ddots & \ddots & \ddots & \end{bmatrix} \]

\[ \small \Qmat_3 = \Jmat_0^T\Qmat_0\Jmat_0 + \sum_{i=1}^{T} (\Jmat_i^T\Qmat_w\Jmat_i) + \Qmat_\beta \]

• Conditioned on the parameters, this is the simplest model possible.
• \( \Qmat \) is large. Example: Antarctica, 4 ST fields, 7 time points, 500k observations, 500k vertices.
• Updated \( \Qmat^* \approx 3.5Gb \). \( \Qmat^{*-1} \approx 1Tb \)
• What if we want to find \( diag(\Amat \Sigma^* \Amat^T) \) where \( \Sigma^* = \Qmat^{*-1} \)?
• Impossible to carry out on a standard PC for large, \( \Amat \), even if \( chol(\Qmat^*) \) is known and fill-in reduced.

• A full backsolve for \( \Amat \Sigma^* \Amat^T \) is wasteful since \( \Amat \) is very sparse.
• We only need a few elements of \( \Sigma^* \): \( [\Amat^T \Amat]_{ij} \ne 0 \)
• Consider the following model:

\[ \begin{align} \yvec &= \Cmat \xvec + \vvec \\ \vvec &\sim \mathcal{N}(0,\Tmat^{-1}) \\ \xvec &\sim \mathcal{N}(0,\Qmat^{-1}) \end{align} \]

• \( prec(\xvec | \yvec) \equiv \Qmat^* = (\Qmat + \Cmat^T \Tmat^{-1} \Cmat) \)
• If \( zeros(\Amat) \subseteq zeros(\Cmat) \) and \( \Tmat \) is diagonal, then \( \Qmat^* \) contains the required elements.

1. In practice the Takahashi equations are always used with GMRFs in order to obtain the marginal variances. Theorem 2 provides a means to obtain uncertainty on linear combinations essentially for free.

2. I can provide predictive uncertainties at any point provided that each row of \( zeros(\Amat) \subseteq zeros(\Cmat) \).

• The updated precision matrix \( \Qmat^* = \Qmat + \Amat^T\Tmat^{-1}\Amat \) is sparse if both \( \Qmat \) and \( \Amat \) are sparse and \( \Tmat \) is diagonal.

Lemma 2: If \( \Amat \) is non-negative then \( [\Amat^T\Amat]_{jk} = 0 \Leftrightarrow diag(\Amat\Sigma^*\Amat^T) \) does not depend on \( \Sigma^*_{jk} \)

• Let \( zeros(\cdot) \) denote the symbolic pattern of a matrix.

\[ \begin{align} zeros(\Qmat^*) &= zeros(\Qmat + \Amat^T\Tmat^{-1}\Amat) \\ &\ge zeros(\Amat^T\Tmat^{-1}\Amat) = zeros(\Amat^T\Amat) \end{align} \]

• In all practical cases there are unknown parameters and a massive GMRF, sampling may be required
• MCMC with multiple chains is relatively straightforward.
  • e.g. Evolutionary MCMC, Embarassingly parallel MCMC.
• Extremely little work on parallelising a single chain.
• Take advantage of local Markov property

\[ x_i \perp \xvec_{-\{i,ne(i)\}} | \xvec_{ne(i)} ~~~ \forall i \in \mathcal{V} \]

\[ E(x_i | \xvec_{-i}) = -\frac{1}{Q_{ii}}\sum_{j:j\sim i}Q_{ij}x_j; ~~~ prec(x_i | \xvec_{-i}) = Q_{ii} \]

• Single-site updating in spatio-temporal systems is a bad idea.
• We can make use of the global Markov property

\[ \xvec_A \perp \xvec_B | \xvec_{C} \]

\( \forall A,B,C \notin \varnothing \) and \( C \) separating \( A,B \).

• In typical Bayesian networks, finding \( A, B \) and \( C \) can be difficult (seed and spawn idea, Gonzales, 2011).

• With a spatial network, this is remarkably easy using the Kernighan-Lin algorithm for graph partitioning (Kernighan and Lin, 1970).